# -*- coding: utf-8 -*-
import logging
import os
import maya.cmds as mc
from liberLibs.dccLibs import Maya
from liberHook.BaseHook import BaseHook
from liberLibs.dcc import selectedNodes
from liberLibs.utils import copytree, write_json
from liberLibs.dccLibs import Xgen


class Hook(BaseHook):
    def __init__(self, library, directory, ext, start, end):
        super(Hook, self).__init__(library, directory, ext, start, end)

    def execute(self):
        # save file
        Maya.save_file()

        # get selected and sort them
        selected_nodes = selectedNodes()
        if len(selected_nodes) != 2:
            self.append_error("Only two nodes supported.")
            return
        if mc.nodeType(selected_nodes[0]) == "xgmPalette":
            selected_nodes.reverse()
        scalp, palette = selected_nodes
        palette = str(palette)
        if mc.nodeType(scalp) != "transform" or mc.nodeType(palette) != "xgmPalette":
            self.append_error("Node selected wrong")
            return
        # copy xgen dir to library/ xgen
        root_project = mc.workspace(q=True, rd=True)
        src_xgen_dir = "%sxgen/collections/%s" % (root_project, palette)
        src_xgen_dir = str(src_xgen_dir.replace("\\", "/"))
        dst_xgen_dir = "%s/xgen_package/collections/%s" % (self.directory, palette)
        dst_xgen_dir = str(dst_xgen_dir.replace("\\", "/"))
        parent_dir = os.path.dirname(dst_xgen_dir)
        if not os.path.isdir(parent_dir):
            os.makedirs(parent_dir)
        try:
            copytree(src_xgen_dir, dst_xgen_dir)
            logging.info("Copy xgen directory done.")
        except Exception as e:
            self.append_error("Copy xgen directory failed.")
            self.append_error(str(e))
            return

        try:
            # set abs path
            x = Xgen.Xgen()
            x.set_abs_path(str(os.path.dirname(parent_dir)))
            logging.info("set abs done.")

            # export scalp abc == abc/.abc
            hair_cache_path = "%s/abc/scalp.abc" % self.directory
            Maya.export_abc(1, 1, hair_cache_path, scalp)
            logging.info("export scalp cache done.")

            # export palette == hair/.xgen
            hair_path = "%s/%s.xgen" % (self.location, self.name)
            hair_path = str(hair_path.replace("\\", "/"))
            x.export_palette(palette, hair_path)
            logging.info("Export .xgen file done.")

            # export shader == shader/.ma
            sg_nodes, connection_dict = self.get_all_hair_sg_nodes(palette)
            if not sg_nodes:
                return
            shd_path = "%s/shader/shader.ma" % self.directory
            shd_path = shd_path.replace("\\", "/")
            connection_path = "%s/shader/shader.json" % self.directory
            connection_path = connection_path.replace("\\", "/")
            mc.select(sg_nodes, r=1, ne=1)
            Maya.export_selected(shd_path)
            write_json(connection_path, connection_dict)
            logging.info("Export shader done.")
            mc.file(new=1, f=1)
        except:
            mc.file(new=1, f=1)

    @staticmethod
    def get_descriptions(palette):
        import xgenm as xgen
        descriptions = xgen.descriptions(palette)
        return descriptions

    def get_all_hair_sg_nodes(self, palette):
        hair_sg_nodes = list()
        connection_dict = dict()
        descriptions = self.get_descriptions(palette)
        if not descriptions:
            return
        for description in descriptions:
            shapes = mc.listRelatives(description, s=1)
            if not shapes:
                continue
            shape = shapes[0]
            sg_nodes = mc.listConnections(shape, s=0, d=1, type="shadingEngine")
            if not sg_nodes:
                continue
            sg_node = sg_nodes[0]
            hair_sg_nodes.append(sg_node)
            connection_dict[description] = sg_node
        return hair_sg_nodes, connection_dict
